import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from scipy.linalg import expm

# 参数设置
n_nodes = 5               # 压力节点数
sigma0 = np.eye(n_nodes) * 0.05  # 目标协方差
dt = 0.1                  # 时间步长
T = 30                    # 总时长
theta_init = np.pi/6      # 初始角度

# PID增益 (对角矩阵)
K_p = np.eye(n_nodes) * 0.5
K_i = np.eye(n_nodes) * 0.1
K_d = np.eye(n_nodes) * 0.2

# 初始化
theta = theta_init
integral_E = np.zeros((n_nodes, n_nodes))
prev_E = np.zeros((n_nodes, n_nodes))

# 模拟压力分布 (简化为角度函数)
def get_pressure(theta):
    noise = np.random.normal(0, 0.01, n_nodes)
    return np.sin(np.linspace(0, np.pi, n_nodes) + theta) + 0.5 + noise

# 计算协方差矩阵
def compute_sigma(p):
    p_mean = np.mean(p, axis=0)
    p_centered = p - p_mean
    return np.cov(p_centered.T)

# 动画初始化
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
ax1.set_xlim(-1, 1)
ax1.set_ylim(0, 2)
ax1.set_title('Seat Pressure Distribution')
nodes = ax1.scatter([], [], c='red', s=100)
angle_text = ax1.text(0.6, 1.8, '', fontsize=10)

ax2.set_title('Covariance Error $\|E\|_F$')
ax2.set_xlim(0, T)
ax2.set_ylim(0, 0.5)
time_line, = ax2.plot([], [], 'r-')
error_text = ax2.text(T*0.7, 0.4, '', fontsize=10)

# 数据记录
time_points = []
error_norms = []

def init():
    nodes.set_offsets(np.zeros((n_nodes, 2)))
    return nodes, time_line

def update(t):
    global theta, integral_E, prev_E
    
    # 1. 获取当前压力并计算协方差
    p = get_pressure(theta)
    sigma_p = compute_sigma(p.reshape(1, -1))  # 需多次采样简化
    
    # 2. 计算协方差误差
    E = sigma_p - sigma0
    error_norm = np.linalg.norm(E, 'fro')
    
    # 3. PID控制 (矩阵形式)
    integral_E += E * dt
    derivative_E = (E - prev_E) / dt
    u = K_p @ E + K_i @ integral_E + K_d @ derivative_E
    
    # 4. 更新角度 (控制信号作用于主关节)
    theta -= np.trace(u) * 0.01  # 标量化控制
    theta = np.clip(theta, 0, np.pi/3)
    prev_E = E
    
    # 5. 更新绘图
    x = np.linspace(-0.8, 0.8, n_nodes)
    y = 0.5 + np.abs(p) * 0.8
    nodes.set_offsets(np.column_stack([x, y]))
    angle_text.set_text(f'Angle: {theta:.2f} rad')
    
    time_points.append(t)
    error_norms.append(error_norm)
    time_line.set_data(time_points, error_norms)
    error_text.set_text(f'Error: {error_norm:.4f}')
    
    return nodes, time_line, angle_text, error_text

ani = FuncAnimation(fig, update, frames=np.arange(0, T, dt),
                    init_func=init, blit=True, interval=100)
ani.save('covariance_pid_control.gif', writer='pillow', dpi=120)
plt.close()